#include <benchmark/benchmark.h>

#include <array>
#include <cstddef>
#include <random>
#include <vector>

// --------------------------------------------------------------
// 常量设定：矩阵的行列数固定为 128×128
// --------------------------------------------------------------
constexpr std::size_t Rows = 128;
constexpr std::size_t Cols = 128;
constexpr std::size_t Total = Rows * Cols;

// --------------------------------------------------------------
// 使用 std::array 构建二维数组别名，行列尺寸在编译期已知
// --------------------------------------------------------------
using Array2D = std::array<std::array<float, Cols>, Rows>;

// --------------------------------------------------------------
// 生成随机浮点数据：返回长度为 Total 的扁平 std::vector<float>
// 每个基准都调用它来初始化输入数据
// --------------------------------------------------------------
std::vector<float> make_data(unsigned seed) {
    std::vector<float> data(Total);
    std::mt19937 rng(seed);  // Mersenne Twister 伪随机数引擎
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    for (auto& v : data) {
        v = dist(rng);
    }
    return data;
}

// --------------------------------------------------------------
// 将扁平 vector 数据拷贝到 C 风格的二维数组中
// - 参数 dst 使用引用避免退化为指针，并保留尺寸信息
// --------------------------------------------------------------
void copy_to_static(float (&dst)[Rows][Cols], const std::vector<float>& src) {
    for (std::size_t i = 0; i < Rows; ++i) {
        for (std::size_t j = 0; j < Cols; ++j) {
            dst[i][j] = src[i * Cols + j];
        }
    }
}

// --------------------------------------------------------------
// 将扁平 vector 数据拷贝到 std::array<std::array<...>> 中
// --------------------------------------------------------------
void copy_to_array(Array2D& dst, const std::vector<float>& src) {
    for (std::size_t i = 0; i < Rows; ++i) {
        for (std::size_t j = 0; j < Cols; ++j) {
            dst[i][j] = src[i * Cols + j];
        }
    }
}

// --------------------------------------------------------------
// 根据扁平 vector 构造 std::vector<std::vector<float>>
// 保证每一行独立分配，模拟常见的“向量的向量”二维结构
// --------------------------------------------------------------
std::vector<std::vector<float>> make_nested(const std::vector<float>& src) {
    std::vector<std::vector<float>> nested(Rows, std::vector<float>(Cols));
    for (std::size_t i = 0; i < Rows; ++i) {
        std::copy_n(src.data() + i * Cols, Cols, nested[i].begin());
    }
    return nested;
}

// --------------------------------------------------------------
// 各种加法核函数：将矩阵 B 加到矩阵 A 上
// --------------------------------------------------------------

// C 风格二维数组版本
void add_static(float (&A)[Rows][Cols], const float (&B)[Rows][Cols]) {
    for (std::size_t i = 0; i < Rows; ++i) {
        for (std::size_t j = 0; j < Cols; ++j) {
            A[i][j] += B[i][j];
        }
    }
}

// std::array<std::array<...>> 版本
void add_array(Array2D& A, const Array2D& B) {
    for (std::size_t i = 0; i < Rows; ++i) {
        for (std::size_t j = 0; j < Cols; ++j) {
            A[i][j] += B[i][j];
        }
    }
}

// 扁平 std::vector 版本：手动计算行偏移
void add_flat(std::vector<float>& A, const std::vector<float>& B) {
    for (std::size_t i = 0; i < Rows; ++i) {
        const std::size_t base = i * Cols;
        for (std::size_t j = 0; j < Cols; ++j) {
            A[base + j] += B[base + j];
        }
    }
}

// 嵌套 std::vector<std::vector<>> 版本：每一行单独访问
void add_nested(std::vector<std::vector<float>>& A,
                const std::vector<std::vector<float>>& B) {
    for (std::size_t i = 0; i < Rows; ++i) {
        auto& rowA = A[i];
        const auto& rowB = B[i];
        for (std::size_t j = 0; j < Cols; ++j) {
            rowA[j] += rowB[j];
        }
    }
}

// --------------------------------------------------------------
// Google Benchmark 基准定义
// 每个基准都会：
//   1. 准备输入数据（只在测量循环外执行）
//   2. 在 state 循环中调用 add_* 内核
//   3. 使用 DoNotOptimize 防止编译器将结果优化掉
//   4. 使用 SetItemsProcessed 记录处理的元素数量
// --------------------------------------------------------------

static void BM_StaticArray(benchmark::State& state) {
    // alignas(64) 保证对齐有利于 SIMD
    alignas(64) float A[Rows][Cols];
    alignas(64) float B[Rows][Cols];
    auto initA = make_data(0x1234u);
    auto initB = make_data(0x5678u);
    copy_to_static(A, initA);
    copy_to_static(B, initB);

    for (auto _ : state) {
        add_static(A, B);
        benchmark::DoNotOptimize(A);
    }
    state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) *
                            static_cast<int64_t>(Total));
}

static void BM_Array2D(benchmark::State& state) {
    Array2D A{};
    Array2D B{};
    auto initA = make_data(0x3344u);
    auto initB = make_data(0x7788u);
    copy_to_array(A, initA);
    copy_to_array(B, initB);

    for (auto _ : state) {
        add_array(A, B);
        benchmark::DoNotOptimize(A);
    }
    state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) *
                            static_cast<int64_t>(Total));
}

static void BM_FlatVector(benchmark::State& state) {
    std::vector<float> A = make_data(0x2468u);
    std::vector<float> B = make_data(0x1357u);

    for (auto _ : state) {
        add_flat(A, B);
        benchmark::DoNotOptimize(A.data());
    }
    state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) *
                            static_cast<int64_t>(Total));
}

static void BM_NestedVector(benchmark::State& state) {
    auto initA = make_data(0x9999u);
    auto initB = make_data(0xAAAAu);
    std::vector<std::vector<float>> A = make_nested(initA);
    std::vector<std::vector<float>> B = make_nested(initB);

    for (auto _ : state) {
        add_nested(A, B);
        benchmark::DoNotOptimize(A.data());
    }
    state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) *
                            static_cast<int64_t>(Total));
}

// 将基准注册到 Google Benchmark 框架
BENCHMARK(BM_StaticArray);
BENCHMARK(BM_Array2D);
BENCHMARK(BM_FlatVector);
BENCHMARK(BM_NestedVector);

// 程序入口由 Google Benchmark 提供
BENCHMARK_MAIN();